import heapq
import random
import networkx as nx
import matplotlib.pyplot as plt
from networkx import shortest_path

class MultiClusterQKDNetwork:
    def __init__(self, clustersPath: list, clusterNum: int):
        self.clusterTopoGraph = nx.Graph()
        self.clustersPath = clustersPath
        self.clusterNum = clusterNum
        self.clusters: list[Cluster] = []
        self.uniqueNodeMap = {}
        self.uniqueNodeNum = 0
        self.maxAlpha: float = 0
        self.maxCompareAlpha: float = 0
        self.clusterTopoGraph.add_nodes_from([i for i in range(1, clusterNum + 1)])
        for i in range(clusterNum):
            path = random.choice(clustersPath)
            c = Cluster(i + 1, path)
            self.clusters.append(c)     # 初始化所有的簇
        self.createClusterLink()
        self.colorClusterNode()
        # self.showClusterTopo()
        # self.clusters[0].showCluster()
        # print(123)
    
    def clearLinkLoad(self):
        for i in self.clusters:
            for j in i.linkList:
                j.load = 0
                j.compareLoad = 0
        self.maxAlpha = 0
        self.maxCompareAlpha = 0
    
    def requestFlow(
            self, 
            startClusterID: int, 
            startNodeID: int, 
            endClusterID: int, 
            endNodeID: int, 
            flow: float
        ):
        # 输入：[[簇ID、节点ID]、[簇ID、节点ID]] 流大小
        # 1. 给定[(簇号,节点号), (簇号,节点号)]
        # 2. 不响应同一个节点的路由
        # 3. 对于不同节点，但是有可能在同一个簇的情况下，按照单簇的 Oblivious Routing 方案进行
        # 4. 求解 Oblivious Routing 的方案
        # 考虑到：
        # 1. 求解大拓扑的路由
        # 2. 求解每个簇间选择哪一条虚边（边界点）
        # 3. 连成一整个 Oblivious Routing 方案
        sp: int = shortest_path(self.clusterTopoGraph, startClusterID, endClusterID)
        # addFlowOnNetwork(self, startID: int, endID: int, flow: float)
        # [1, 103]   NodeID -> NodeID_1
        # [103, 172] NodeID_1 -> NodeID_2
        # [172, 199] NodeID_2 -> NodeID_3
        def addFlow(clusterID: int, startID: int, endID: int):
            if(startID == endID):
                return
            else:
                # print(f"增加流clusterID:{clusterID} startID:{startID} endID:{endID}")
                # 每次增加流 就获取一次Alpha的值
                self.clusters[clusterID - 1].addSpfFlowOnNetwork(startID, endID, flow)
                self.clusters[clusterID - 1].addCplexFlowOnNetwork(startID, endID, flow)
                t = self.clusters[clusterID - 1].getAlpha()
                tCompare = self.clusters[clusterID - 1].getCompareAlpha()

                if(t > self.maxAlpha):
                    self.maxAlpha = t
                if(tCompare > self.maxCompareAlpha):
                    self.maxCompareAlpha = tCompare
                    
        # print(sp)

        bridgeNodeID = 0
        
        for i in range(len(sp) - 1):
            cluster_id_1: int = sp[i]
            cluster_id_2: int = sp[i + 1]
            for adjacentInfo in self.clusters[cluster_id_1 - 1].adjacentClusterList:
                if(adjacentInfo.adjacent_cluster_id == cluster_id_2):
                    # 寻找跨簇节点
                    bridgeNodeID = adjacentInfo.node_id
                    addFlow(cluster_id_1, startNodeID, bridgeNodeID)
                    startNodeID = adjacentInfo.adjacent_node_id
                    break

        addFlow(sp[-1], startNodeID, endNodeID)
        return

    def showClusterTopo(self):
        nx.draw(
            self.clusterTopoGraph, 
            with_labels=True, 
            node_color='red', 
            node_size=200, 
            font_size=10, 
            font_weight='bold'
        )
        plt.show()
    
    def addEdgeToTopoGraph(self, edge):
        # 构建簇的整体拓扑
        tmp_g = self.clusterTopoGraph
        if not tmp_g.has_edge(*edge) and not tmp_g.has_edge(edge[1], edge[0]) and edge[0] != edge[1]:
            self.clusterTopoGraph.add_edge(*edge)
    
    def getAdjacentNodeColor(self, cluster_id: int, node_id: int) -> int:
        # 判断该节点相邻簇的相连点是否有已经染色的节点，若有，返回这个节点uniqueID，若无，返回0
        ret_unique_id = 0
        # repeat_count = 0
        for adjacent in self.clusters[cluster_id - 1].adjacentClusterList:
            if(adjacent.node_id == node_id):
                t = self.clusters[adjacent.adjacent_cluster_id - 1].nodeList[adjacent.adjacent_node_id - 1].uniqueID
                if(t != 0):
                    # repeat_count += 1
                    ret_unique_id = t
                    # if(repeat_count != 1 and repeat_count != 0):
                        # print(f"重复：{repeat_count} cluster_id:{cluster_id} node_id:{node_id} adjacent_cluster_id:{adjacent.adjacent_cluster_id} adjacent_node_id:{adjacent.adjacent_node_id}")
        return ret_unique_id

    def colorClusterNode(self):
        # 为大簇中的每个节点，染上唯一的颜色（uniqueID）
        count = 1
        for i in range(len(self.clusters)):
            cluster = self.clusters[i]
            for j in range(len(cluster.nodeList)):
                node = cluster.nodeList[j]
                if(node.uniqueID == 0):
                    tmp_color = self.getAdjacentNodeColor(cluster.cluster_id, node.nodeID)
                    if(tmp_color == 0):
                        self.clusters[i].nodeList[j].uniqueID = count #染色
                    else:
                        self.clusters[i].nodeList[j].uniqueID = tmp_color #染色
                    if(node.uniqueID in self.uniqueNodeMap):
                        self.uniqueNodeMap[node.uniqueID] += 1
                    else:
                        self.uniqueNodeMap[node.uniqueID] = 1
                        self.uniqueNodeNum += 1
                    for k in range(len(cluster.adjacentClusterList)):
                        adjacentCluster = cluster.adjacentClusterList[k]
                        adjacent_cluster_id = adjacentCluster.adjacent_cluster_id
                        adjacent_node_id = adjacentCluster.adjacent_node_id
                        if(adjacentCluster.node_id == node.nodeID):
                            # 需要把相邻簇的节点，也染成同样的颜色
                            self.clusters[adjacent_cluster_id - 1].nodeList[adjacent_node_id - 1].uniqueID = count
                    count += 1
    
    def createClusterLink(self):
        # 每两个簇之间，随机创建3条链路
        for i in range(self.clusterNum):
            cluster = self.clusters[i] # 取出每个簇
            for _ in range(2): # 每个簇找其它两个簇，连3条边
                cluster_tmp_index = random.choice([elem for j, elem in enumerate(range(self.clusterNum)) if j != i])
                cluster_tmp = self.clusters[cluster_tmp_index]
                adjacent_cluster_id = cluster_tmp.cluster_id #邻居簇的ID
                adjacent_cluster_node_num = cluster_tmp.nodeNum
                for _ in range(3):
                    node_id = random.randint(1, cluster.nodeNum) #随机一个簇内节点ID
                    adjacent_node_id = random.randint(1, adjacent_cluster_node_num) #随机一个邻居簇节点的ID
                    self.clusters[i].adjacentClusterList.append( #把他们连起来
                        Cluster.adjacentCluster(
                            adjacent_cluster_id,
                            node_id,
                            adjacent_node_id
                        )
                    )
                    self.clusters[adjacent_cluster_id - 1].adjacentClusterList.append( #把他们连起来
                        Cluster.adjacentCluster(
                            i + 1,
                            adjacent_node_id,
                            node_id
                        )
                    )
                    self.addEdgeToTopoGraph((i + 1, adjacent_cluster_id))

    def getClusterTopo():
        print("获取簇的拓扑，仅仅是簇的而已，用于整个网络的dfs")

class Cluster:
    def __init__(self, cluster_id: int , path):
        self.cluster_id = cluster_id
        self.txt_path = path
        self.alpha: float = 0
        self.beta: float = 0
        self.seed: int = 0
        self.nodeNum: int = 0
        self.nodeList: list[Cluster.Node] = []
        self.linkList: list[Cluster.Link] = []
        self.obliviousPlanList: list[Cluster.ObliviousPlan] = []
        self.adjacentClusterList: list[Cluster.adjacentCluster] = [] # 暂时置空，需要的时候再连边
        self.readDataFromTxt(path)
    
    class Link:
        def __init__(self, startNodeID: int, endNodeID: int, capacity: int, load: float  = 0, compareLoad: float = 0):
            self.startNodeID = startNodeID
            self.endNodeID = endNodeID
            self.capacity = capacity
            self.load = load
            self.compareLoad = compareLoad

    class Node:
        def __init__(self, nodeID):
            self.nodeID = nodeID        # 节点ID
            self.uniqueID = 0
            self.adjacentNodes = []     # 相邻节点ID
            self.capacity = 0           # 节点负载量：节点相邻的链路容量之和
    
    class PathLoad:
        # 用户判断源目节点对之间Oblivious路径的复杂情况
        def __init__(self, m: int, n: int, load_index: float):
            self.m = m
            self.n = n
            self.load_index = load_index

    class ObliviousPlan:
        class fixFlow:
            def __init__(self, link: list[int], ratio: float):
                self.link = link
                self.ratio = ratio
        def __init__(self, startID: int, endID: int, route: list, fixPath: list): # Oblivious方案
            self.startID = startID
            self.endID = endID
            self.route = route
            # [
            #  { link: [i, j, a, b],  ratio: 0.3442 }, 
            #  { link: [i, j, a, b],  ratio: 0.3442 }, 
            # ]
            self.fixPath: list[Cluster.ObliviousPlan.fixFlow] = fixPath
        
    class adjacentCluster:
        # 相邻簇节点 簇ID 节点ID 簇连点ID
        def __init__(self, adjacent_cluster_id: int, node_id: int, adjacent_node_id: int):
            self.adjacent_cluster_id = adjacent_cluster_id
            self.node_id = node_id
            self.adjacent_node_id = adjacent_node_id

    def showCluster(self):
        with open(self.txt_path, 'r') as file:
            lines = file.readlines()
        n = int(lines[0])
        G = nx.Graph()
        G.add_nodes_from(range(1, n + 1))
        edge_nums = int(lines[2])
        for i in range(3, 3 + edge_nums):
            start, end, weight = map(int, lines[i].split())
            G.add_edge(start, end, weight=weight)
        fig, ax = plt.subplots(figsize=(50, 50))
        pos = nx.spring_layout(G, seed=41)
        nx.draw(G, pos, ax=ax, with_labels=True, node_color='red', node_size=300, font_size=8, alpha=1)
        edge_labels = nx.get_edge_attributes(G, 'weight')
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=6)
        plt.show()
    
    def getCompareAlpha(self) -> float:
        # 获取网络负载系数 max { Link CompareLoad / Link Capacity }
        alpha = 0
        for link in self.linkList:
            if(link.compareLoad / link.capacity > alpha):
                alpha = link.compareLoad / link.capacity
        return alpha
    
    def getAlpha(self) -> float:
        # 获取网络负载系数 max { Link Load / Link Capacity }
        alpha = 0
        for link in self.linkList:
            if(link.load / link.capacity > alpha):
                alpha = link.load / link.capacity
        return alpha

    def addSpfFlowOnNetwork(self, startID: int, endID: int, flow: float):
        # Dijkstra算法查找最短路径
        dist = {node.nodeID: float('inf') for node in self.nodeList}
        dist[startID] = 0
        prev_node = {node.nodeID: None for node in self.nodeList}
        pq = [(0, startID)]
        while pq:
            curr_dist, curr_node = heapq.heappop(pq)
            if curr_node == endID:
                break
            if curr_dist > dist[curr_node]:
                continue
            # 更新相邻节点的距离
            for neighbor_id in self.getNode(curr_node).adjacentNodes:
                new_dist = curr_dist + 1
                if new_dist < dist[neighbor_id]:
                    dist[neighbor_id] = new_dist
                    prev_node[neighbor_id] = curr_node
                    heapq.heappush(pq, (new_dist, neighbor_id))

        # 输出最短路径
        path = []
        node = endID
        while node is not None:
            path.insert(0, node)
            node = prev_node[node]

        for i in range(len(path) - 1):
            # 增加流
            _start = path[i]
            _end = path[i + 1]
            self.addCompareLinkLoad(_start, _end, flow)

    def addCplexFlowOnNetwork(self, startID: int, endID: int, flow: float):
        # 按照 Oblivious 方案，为网络施加一个流 
        # 输入：起点、终点
        # 操作：为对应的边添加流量
        for plan in self.obliviousPlanList:
            if(plan.startID == startID and plan.endID == endID):
                for fix in plan.fixPath:
                    self.addLinkLoad(fix.link[2], fix.link[3], flow * fix.ratio)

    def addCompareLinkLoad(self, startID, endID, load):
        # 设置链路负载，主要是针对对比的流进行
        for link in self.linkList:
            if link.startNodeID == startID and link.endNodeID == endID:
                link.compareLoad += load
            if link.startNodeID == endID and link.endNodeID == startID:
                link.compareLoad += load

    def addLinkLoad(self, startID, endID, load):
        # 设置链路负载，主要是针对Oblivious Routing的流进行
        for link in self.linkList:
            if link.startNodeID == startID and link.endNodeID == endID:
                link.load += load
            if link.startNodeID == endID and link.endNodeID == startID:
                link.load += load

    def getLinkCapacity(self, startID, endID):
        for link in self.linkList:
            if link.startNodeID == startID and link.endNodeID == endID:
                return link.capacity
            if link.startNodeID == endID and link.endNodeID == startID:
                return link.capacity
        return -1

    def getNode(self, nodeId) -> Node:
        if nodeId - 1 >= len(self.nodeList) or nodeId - 1 < 0:
            return self.Node(-1)
        return self.nodeList[nodeId - 1]

    def readDataFromTxt(self, path):
        # 从txt文件中读入网络簇
        with open(path, 'r') as file:
            lines = file.readlines()
            numNodes = int(lines[0]) #节点数量
            self.nodeNum = numNodes
            nodeIDs = [i for i in range(1, numNodes + 1)] #节点ID
            for nodeID in nodeIDs: #将节点添加至节点列表中
                self.nodeList.append(self.Node(nodeID))
            param = lines[1].split() #簇的参数
            self.alpha = float(param[0])
            self.beta = float(param[1])
            self.seed = int(param[2])
            numLinks = int(lines[2]) #边数量

            # 初始化边以及边容量
            for i in range(3, 3 + numLinks):
                linkData = lines[i].split()
                startNode = int(linkData[0])
                endNode = int(linkData[1])
                self.nodeList[startNode - 1].adjacentNodes.append(endNode)
                self.nodeList[endNode - 1].adjacentNodes.append(startNode)
                capacity = int(linkData[2])
                link = self.Link(startNode, endNode, capacity)
                self.linkList.append(link)
            
            sd_num = numNodes * (numNodes - 1)
            cursor = 3 + numLinks
            for i in range(sd_num):
                route = lines[cursor].split() #路径长度-源节点1-目的节点1
                routeLength = int(route[0]) #路径长度
                startID = int(route[1])
                endID = int(route[2])
                routePath = []
                fixPath: list[Cluster.ObliviousPlan.fixFlow] = []
                for _ in range(routeLength):
                    cursor = cursor + 1
                    fix_line = lines[cursor].split() #边点1-边点2-点a-点b-fixRatio-点c-点d-fixRatio
                    link_start_node = int(fix_line[0])
                    link_end_node = int(fix_line[1])
                    routePath.append([link_start_node, link_end_node])
                    fix_num = int((len(fix_line) - 2) / 3)
                    for j in range(fix_num):
                        fix_cursor = 2 + j * 3
                        link_i = int(fix_line[fix_cursor])
                        link_j = int(fix_line[fix_cursor + 1])
                        link_ratio = float(fix_line[fix_cursor + 2])
                        fixPath.append(Cluster.ObliviousPlan.fixFlow(
                            [link_start_node, link_end_node, link_i, link_j],
                            link_ratio
                        ))
                # print(routePath)
                cursor = cursor + 1
                tmp = Cluster.ObliviousPlan(
                    startID, 
                    endID, 
                    routePath, 
                    fixPath
                )
                self.obliviousPlanList.append(tmp)

            
